#!/usr/bin/env python3
"""
Quick utility to inspect the positive / negative label distribution in an
exported LWG training CSV.

Example:
    python tools/lwg/count_label_distribution.py \
        --csv outputs/dancetrack/lwg_training/lwg_train.csv
"""

import argparse
import csv
from collections import Counter
from pathlib import Path


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Count label distribution for an exported LWG CSV."
    )
    parser.add_argument(
        "--csv",
        type=str,
        required=True,
        help="Path to the LWG training CSV (e.g. outputs/.../lwg_train.csv).",
    )
    args = parser.parse_args()

    csv_path = Path(args.csv)
    if not csv_path.exists():
        raise FileNotFoundError(f"{csv_path} not found.")

    counter: Counter[str] = Counter()
    with csv_path.open("r") as f:
        reader = csv.DictReader(f)
        for row in reader:
            label = row.get("y")
            if label is None:
                raise KeyError("CSV missing column 'y'.")
            counter[label] += 1

    total = sum(counter.values())
    print(f"Total samples: {total}")
    for label, count in sorted(counter.items()):
        ratio = (count / total * 100.0) if total > 0 else 0.0
        print(f"  y={label}: {count} ({ratio:.2f}%)")


if __name__ == "__main__":
    main()

